import numpy as np
from scipy.linalg import pinv
import matplotlib.pyplot as plt
from tqdm import tqdm
import os


def error(x, pointList, distanceList):
    sum1=0
    for point,dis in zip(pointList,distanceList):
        sum=0
        for upx, px in zip(x.tolist()[0], point):
            sum += (px-upx)**2
        sum1+=abs(sum-dis**2)
    return sum1


class PSOPoint:
    def __init__(self, dimention,domain):
        maxs = max([di[1] for di in domain])
        mins = min([di[0] for di in domain])
        self.x = np.random.rand(1, dimention)*(maxs-mins)+mins
        self.v = np.random.rand(1, dimention)*(maxs-mins)*0.5
        self.myHistoryBestx = self.x
        self.myHistoryBest = 999999999

    def setLocation(self, newLocation):
        self.x = newLocation

    def setvelocity(self, newVelocity):
        self.velocity = newVelocity

    def setHistoryBest(self, newhistoryBest):
        self.myHistoryBest = newhistoryBest

    def setHistoryBestx(self, newhistoryBestx):
        self.myHistoryBestx = newhistoryBestx

class PSO:
    def __init__(self, maxVelocity, minVelocity, locationLimit, maxEpoch, f, dimension, pointCount, Wmax, Wmin, Cmin, Cmax,existedPointList, distanceList, domain, precision):
        self.precision = precision
        self.existedPointList = existedPointList
        self.distanceList = distanceList
        self.max_v = maxVelocity
        self.min_v = minVelocity
        self.max_ep = maxEpoch
        self.f = f
        self.domain = domain
        self.pointList = [PSOPoint(dimension,domain) for i in range(0, pointCount)]
        fx = [f(point.x,self.existedPointList,self.distanceList) for point in self.pointList]
        self.globalbestx = self.pointList[fx.index(min(fx))].x
        self.globalbest = min(fx)
        self.Wmax = Wmax
        self.Wmin = Wmin
        self.Cmin = Cmin
        self.Cmax = Cmax
        self.imgcount = 0
        self.locationLimit = locationLimit
        self.dimension = dimension

    def personalBest(self, point, currentX):
        if point.myHistoryBest < self.f(currentX,self.existedPointList,self.distanceList):
            if self.Annealing(point):
                point.setHistoryBest(self.f(currentX,self.existedPointList,self.distanceList))
                point.setHistoryBestx(currentX)
        else:
            point.setHistoryBest(self.f(currentX,self.existedPointList,self.distanceList))
            point.setHistoryBestx(currentX)

    def globalBest(self):
        for point in self.pointList:
            if self.f(point.x,self.existedPointList,self.distanceList) < self.globalbest:
                self.globalbest = self.f(point.x,self.existedPointList,self.distanceList)
                self.globalbestx = point.x

    def calNewVelocity(self, point, ep):
        r1 = np.random.uniform(0, 1)
        r2 = np.random.uniform(0, 1)
        self.C1 = self.C2 = (self.Cmax-(self.Cmax-self.Cmin)*ep/self.max_ep)
        v = (self.Wmax-(self.Wmax-self.Wmin)*(ep/self.max_ep)**2/self.max_ep)*point.v + \
            self.C1*r1*(point.myHistoryBestx-point.x) + \
            self.C2*r2*(self.globalbestx-point.x)
        v[v > self.max_v] = self.max_v
        v[v < self.min_v] = self.min_v
        return v

    def Annealing(self, point):
        p = np.exp(-point.v/self.max_v)
        p = np.average(p)
        r = np.random.rand(1)
        if r > p:
            return True
        else:
            return False

    def run(self,saveImg):
        oldx = self.pointList[0].x
        for ep in tqdm(range(0, self.max_ep)):
            for index, point in enumerate(self.pointList):
                if self.Annealing(point) or True:
                    newVelocity = self.calNewVelocity(point, ep)
                    self.pointList[index].setvelocity(newVelocity)
                    newLocation = self.pointList[index].x+newVelocity
                    self.pointList[index].setLocation(newLocation)
                self.personalBest(point, point.x)
            self.globalBest()
            if self.globalbest<self.precision*sum([i**2 for i in self.distanceList]):
                break

            newx = self.pointList[0].x
            oldx = newx
            if ep % 1 == 0 and saveImg:
                self.drawPoints()
                pass
        print(self.globalbest)
        print(self.globalbestx)

    def printPointLocation(self):
        for point in self.pointList:
            print(point.x)
        print("=====================================")

    def drawPoints(self):
        plt.figure(figsize=(10,10))
        x=[epoint[0] for epoint in self.existedPointList]
        y=[epoint[1] for epoint in self.existedPointList]
        plt.scatter(x,y,color='purple',marker="o",s=300)
        for epoint,d in zip(self.existedPointList,self.distanceList):
            self.drawCircle(epoint[0],epoint[1],d)
        x_list = []
        y_list = []
        for point in self.pointList:
            x_list.append(point.x[0][0])
            y_list.append(point.x[0][1])
        plt.xlim(self.domain[0])
        plt.ylim(self.domain[1])
        plt.plot(x_list, y_list, marker="o", linestyle="")
        plt.savefig("./PSOimages/%d.png" % self.imgcount)
        self.imgcount += 1
        plt.close()
    
    def drawCircle(self,x0,y0,r):
        theta = np.linspace(0, 2 * np.pi, 200)
        x = np.cos(theta)*r+x0
        y = np.sin(theta)*r+y0
        plt.plot(x, y)
        # plt.show()


def Solve(pointList, distanceList, domain, precision):
    span = [abs(di[0]-di[1]) for di in domain]
    pointNum = np.prod(span)/8
    myPSO = PSO(maxVelocity=min(span)/2,
                minVelocity=-min(span)/2,
                maxEpoch=5000,
                locationLimit=None,
                f=error,
                dimension=len(pointList[0]),
                pointCount=128,
                Wmax=0,
                Wmin=1,
                Cmin=0.5,
                Cmax=0.5,
                existedPointList=pointList, 
                distanceList=distanceList,
                domain=domain,
                precision = precision)
    myPSO.run(saveImg=False)


if __name__ == "__main__":
    # 三个已知点
    ThreePoints = [[2.36, 11.14],
                   [10.3,6.34], 
                   [15.63,13.29]]
    # 到三点的距离
    ThreeDistance = [7.74,6.97, 5.83]
    # 可以肉眼判断出待测点的大致范围
    domain = [[0,20],[0,20]]
    # 运行粒子群
    Solve(pointList=ThreePoints, distanceList=ThreeDistance, domain=domain, precision=0.0005)

